import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.interpolate import interp1d
 
# Step 1: Interpolate centerline to ensure uniform spacing 
def interpolate_centerline(centerline, num_points):
    t = np.linspace(0, 1, len(centerline))
    fx = interp1d(t, centerline[:, 0], kind='cubic')
    fy = interp1d(t, centerline[:, 1], kind='cubic')
    fz = interp1d(t, centerline[:, 2], kind='cubic')
    t_new = np.linspace(0, 1, num_points)
    new_centerline = np.vstack([fx(t_new), fy(t_new), fz(t_new)]).T 
    return new_centerline
 
# Step 2: Generate cross-sections 
def generate_cross_sections(centerline, radii, num_points=16):
    cross_sections = []
    prev_N = None  # Store previous normal vector for consistency
    
    for i in range(len(centerline)):
        # Debugging: Print current point index and radius 
        print(f"Processing point {i} of {len(centerline)}, radius={radii[i]}")
        
        # Compute tangent vector
        if i == 0:
            T = centerline[i + 1] - centerline[i]
        elif i == len(centerline) - 1:
            T = centerline[i] - centerline[i - 1]
        else:
            T = centerline[i + 1] - centerline[i - 1]
        
        T = T / np.linalg.norm(T)
        
        # Compute normal and binormal vectors 
        V = np.array([1, 0, 0]) if T[0] != 1 else np.array([0, 1, 0])
        N = np.cross(T, V)
        N = N / np.linalg.norm(N)
        B = np.cross(T, N)
        
        # Ensure consistent orientation of cross-sections
        if prev_N is not None:
            dot_product = np.dot(N, prev_N)
            if dot_product < 0:
                N = -N
                B = -B 
        prev_N = N
        
        # Generate points on the cross-section with varying radius
        theta = np.linspace(0, 2 * np.pi, num_points, endpoint=False)
        points = centerline[i] + radii[i] * (np.outer(np.cos(theta), N) + np.outer(np.sin(theta), B))
        cross_sections.append(points)
    
    return cross_sections
 
# Step 3: Connect cross-sections to form a mesh
def connect_cross_sections(cross_sections):
    vertices = []
    faces = []
    vertex_offset = 0
    
    for i in range(len(cross_sections) - 1):
        current_section = cross_sections[i]
        next_section = cross_sections[i + 1]
        
        # Add vertices
        vertices.extend(current_section)
        #vertices.extend(next_section)
        
        # Add faces
        num_points = len(current_section)
        for j in range(num_points):
            j_next = (j + 1) % num_points
            faces.append([vertex_offset + j, vertex_offset + j_next, vertex_offset + num_points + j_next])
            faces.append([vertex_offset + j, vertex_offset + num_points + j_next, vertex_offset + num_points + j])
        
        # Debugging: Print vertex offset and face indices 
        print(f"Processing cross-section pair {i}-{i+1}, vertex_offset={vertex_offset}")
        print(f"Faces added: {[face for face in faces[-num_points*2:]]}")
        
        vertex_offset += num_points  # Increment vertex offset correctly 
    vertices.extend(next_section)
    return np.array(vertices), np.array(faces)
 
# Step 4: Visualize the mesh and centerline 
def visualize_mesh_and_centerline(vertices, faces, centerline):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    # Create a collection of triangles for the mesh
    mesh = Poly3DCollection(vertices[faces], alpha=0.5, edgecolor='k', facecolor='cyan')
    ax.add_collection3d(mesh)
    
    # Plot the centerline
    centerline_x = centerline[:, 0]
    centerline_y = centerline[:, 1]
    centerline_z = centerline[:, 2]
    ax.plot(centerline_x, centerline_y, centerline_z, color='red', linewidth=2, label="Centerline")
    
    # Set plot limits
    ax.set_xlim([vertices[:, 0].min(), vertices[:, 0].max()])
    ax.set_ylim([vertices[:, 1].min(), vertices[:, 1].max()])
    ax.set_zlim([vertices[:, 2].min(), vertices[:, 2].max()])
    
    # Add labels and legend 
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.legend()
    
    plt.show()
 
# Example usage
if __name__ == "__main__":
    # Simulated centerline and radii
    centerline = np.array([[0, 0, 0], [1, 0, 0], [2, 1, 0], [3, 1, 1]])
    radii = np.array([0.05, 0.17, 0.06, 0.08])  # Different radii for each point
    
    # Interpolate centerline
    centerline = interpolate_centerline(centerline, num_points=20)
    radii = np.interp(np.linspace(0, 1, len(centerline)), np.linspace(0, 1, len(radii)), radii)  # Interpolate radii
    
    # Generate cross-sections
    cross_sections = generate_cross_sections(centerline, radii)
    
    # Connect cross-sections
    vertices, faces = connect_cross_sections(cross_sections)
    
    # Debugging: Print the total number of vertices and faces
    print(f"Total number of vertices: {len(vertices)}")
    print(f"Total number of faces: {len(faces)}")
    print(f"Maximum vertex index in faces: {np.max(faces)}")
    
    # Visualize the mesh and centerline
    visualize_mesh_and_centerline(vertices, faces, centerline)